package org.apache.flink.runtime.resourceprovider;

import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executor;

import javax.annotation.Nullable;
import javax.annotation.concurrent.GuardedBy;
import javax.annotation.concurrent.ThreadSafe;

import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.api.common.ConsumerID;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.runtime.highavailability.HighAvailabilityServices;
import org.apache.flink.runtime.leaderretrieval.LeaderRetrievalListener;
import org.apache.flink.runtime.leaderretrieval.LeaderRetrievalService;
import org.apache.flink.runtime.providermanager.UnresolvedProviderManagerLocation;
import org.apache.flink.runtime.registration.RegisteredRpcConnection;
import org.apache.flink.runtime.registration.RegistrationResponse;
import org.apache.flink.runtime.registration.RetryingRegistration;
import org.apache.flink.runtime.registration.RetryingRegistrationConfiguration;
import org.apache.flink.runtime.resourceconsumer.RCRPRegistrationRejection;
import org.apache.flink.runtime.resourceconsumer.RCRPRegistrationSuccess;
import org.apache.flink.runtime.resourceconsumer.ResourceConsumerGateway;
import org.apache.flink.runtime.resourceconsumer.ConsumingMasterId;
import org.apache.flink.runtime.resourceconsumer.ResourceProviderRegistrationInformation;
import org.apache.flink.runtime.rpc.RpcService;
import org.apache.flink.util.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Default implementation of {@link ConsumerLeaderService}. */
public class DefaultConsumerLeaderService implements ConsumerLeaderService {

    private static final Logger LOG = LoggerFactory.getLogger(DefaultConsumerLeaderService.class);

    /** Self's location, used for the job manager connection. */
    private final UnresolvedProviderManagerLocation ownLocation;

    /** The leader retrieval service and listener for each registered consumer. */
    private final Map<
                    ConsumerID,
                    Tuple2<LeaderRetrievalService,
                            DefaultConsumerLeaderService.ResourceConsumerLeaderListener>> consumerLeaderServices;

    private final RetryingRegistrationConfiguration retryingRegistrationConfiguration;

    private final UUID taskManagerSession;

    /** Internal state of the service. */
    private volatile DefaultConsumerLeaderService.State state;

    /**
     * Address of the owner of this service. 
     * This address is used for the resource consumer connection.
     */
    private String ownerAddress;

    /** Rpc service to use for establishing connections. */
    private RpcService rpcService;

    /** High availability services to create the leader retrieval services from. */
    private HighAvailabilityServices highAvailabilityServices;

    /** Job leader listener listening for job leader changes. */
    private ConsumerLeaderListener consumerLeaderListener;

    public DefaultConsumerLeaderService(
            UnresolvedProviderManagerLocation location,
            RetryingRegistrationConfiguration retryingRegistrationConfiguration) {
        this.ownLocation = Preconditions.checkNotNull(location);
        this.retryingRegistrationConfiguration =
                Preconditions.checkNotNull(retryingRegistrationConfiguration);
        this.taskManagerSession = UUID.randomUUID();

        // Has to be a concurrent hash map because tests might access this service
        // concurrently via containsJob
        consumerLeaderServices = new ConcurrentHashMap<>(4);

        state = DefaultConsumerLeaderService.State.CREATED;

        ownerAddress = null;
        rpcService = null;
        highAvailabilityServices = null;
        consumerLeaderListener = null;
    }

    // -------------------------------------------------------------------------------
    // Methods
    // -------------------------------------------------------------------------------

    @Override
    public void start(
            final String initialOwnerAddress,
            final RpcService initialRpcService,
            final HighAvailabilityServices initialHighAvailabilityServices,
            final ConsumerLeaderListener initialJobLeaderListener) {

        if (DefaultConsumerLeaderService.State.CREATED != state) {
            throw new IllegalStateException("The service has already been started.");
        } else {
            LOG.info("Start job leader service.");

            this.ownerAddress = Preconditions.checkNotNull(initialOwnerAddress);
            this.rpcService = Preconditions.checkNotNull(initialRpcService);
            this.highAvailabilityServices =
                    Preconditions.checkNotNull(initialHighAvailabilityServices);
            this.consumerLeaderListener = Preconditions.checkNotNull(initialJobLeaderListener);
            state = DefaultConsumerLeaderService.State.STARTED;
        }
    }

    @Override
    public void stop() throws Exception {
        LOG.info("Stop job leader service.");

        if (DefaultConsumerLeaderService.State.STARTED == state) {

            for (Tuple2<LeaderRetrievalService, DefaultConsumerLeaderService.ResourceConsumerLeaderListener>
                    leaderRetrievalServiceEntry : consumerLeaderServices.values()) {
                LeaderRetrievalService leaderRetrievalService = leaderRetrievalServiceEntry.f0;
                DefaultConsumerLeaderService.ResourceConsumerLeaderListener jobManagerLeaderListener =
                        leaderRetrievalServiceEntry.f1;

                jobManagerLeaderListener.stop();
                leaderRetrievalService.stop();
            }

            consumerLeaderServices.clear();
        }

        state = DefaultConsumerLeaderService.State.STOPPED;
    }

    @Override
    public void removeJob(ConsumerID jobId) {
        Preconditions.checkState(
                DefaultConsumerLeaderService.State.STARTED == state,
                "The service is currently not running.");

        Tuple2<LeaderRetrievalService, DefaultConsumerLeaderService.ResourceConsumerLeaderListener> entry =
                consumerLeaderServices.remove(jobId);

        if (entry != null) {
            LOG.info("Remove job {} from job leader monitoring.", jobId);

            LeaderRetrievalService leaderRetrievalService = entry.f0;
            DefaultConsumerLeaderService.ResourceConsumerLeaderListener jobManagerLeaderListener = entry.f1;

            jobManagerLeaderListener.stop();

            try {
                leaderRetrievalService.stop();
            } catch (Exception e) {
                LOG.info(
                        "Could not properly stop the LeaderRetrievalService for job {}.", jobId, e);
            }
        }
    }

    @Override
    public void addJob(final ConsumerID consumerId, final String defaultTargetAddress) throws Exception {
        Preconditions.checkState(
                DefaultConsumerLeaderService.State.STARTED == state,
                "The service is currently not running.");

        LOG.info("Add job {} for job leader monitoring.", consumerId);

        final LeaderRetrievalService leaderRetrievalService =
                highAvailabilityServices.getResourceConsumerLeaderRetriever(defaultTargetAddress);

        DefaultConsumerLeaderService.ResourceConsumerLeaderListener jobManagerLeaderListener =
                new ResourceConsumerLeaderListener(consumerId);

        final Tuple2<LeaderRetrievalService, ResourceConsumerLeaderListener> oldEntry =
                consumerLeaderServices.put(
                        consumerId, Tuple2.of(leaderRetrievalService, jobManagerLeaderListener));

        if (oldEntry != null) {
            oldEntry.f0.stop();
            oldEntry.f1.stop();
        }

        leaderRetrievalService.start(jobManagerLeaderListener);
    }

    @Override
    public void reconnect(final ConsumerID jobId) {
        Preconditions.checkNotNull(jobId, "JobID must not be null.");

        final Tuple2<LeaderRetrievalService, ResourceConsumerLeaderListener> jobLeaderService =
                consumerLeaderServices.get(jobId);

        if (jobLeaderService != null) {
            jobLeaderService.f1.reconnect();
        } else {
            LOG.info("Cannot reconnect to job {} because it is not registered.", jobId);
        }
    }

    /** Leader listener which tries to establish a connection to a newly detected consumer leader. */
    @ThreadSafe
    private final class ResourceConsumerLeaderListener implements LeaderRetrievalListener {

        private final Object lock = new Object();

        /** Consumer id identifying the consumer to look for a leader. */
        private final ConsumerID consumerId;

        /** Rpc connection to the job leader. */
        @GuardedBy("lock")
        @Nullable
        private RegisteredRpcConnection<
        				ConsumingMasterId,
        				ResourceConsumerGateway,
                        RCRPRegistrationSuccess,
                        RCRPRegistrationRejection>
                rpcConnection;

        /** Leader id of the current resource consumer leader. */
        @GuardedBy("lock")
        @Nullable
        private ConsumingMasterId currentResourceConsumerId;

        /** State of the listener. */
        private volatile boolean stopped;

        private ResourceConsumerLeaderListener(ConsumerID jobId) {
            this.consumerId = Preconditions.checkNotNull(jobId);

            stopped = false;
            rpcConnection = null;
            currentResourceConsumerId = null;
        }

        private ConsumingMasterId getCurrentJobMasterId() {
            synchronized (lock) {
                return currentResourceConsumerId;
            }
        }

        public void stop() {
            synchronized (lock) {
                if (!stopped) {
                    stopped = true;

                    closeRpcConnection();
                }
            }
        }

        public void reconnect() {
            synchronized (lock) {
                if (stopped) {
                    LOG.debug(
                            "Cannot reconnect because the JobManagerLeaderListener has already been stopped.");
                } else {
                    if (rpcConnection != null) {
                        Preconditions.checkState(
                                rpcConnection.tryReconnect(),
                                "Illegal concurrent modification of the JobManagerLeaderListener rpc connection.");
                    } else {
                        LOG.debug("Cannot reconnect to an unknown JobMaster.");
                    }
                }
            }
        }

        @Override
        public void notifyLeaderAddress(
                @Nullable final String leaderAddress, @Nullable final UUID leaderId) {
            Optional<ConsumingMasterId> jobManagerLostLeadership = Optional.empty();

            synchronized (lock) {
                if (stopped) {
                    LOG.debug(
                            "{}'s leader retrieval listener reported a new leader for job {}. "
                                    + "However, the service is no longer running.",
                            DefaultConsumerLeaderService.class.getSimpleName(),
                            consumerId);
                } else {
                    final ConsumingMasterId jobMasterId = ConsumingMasterId.fromUuidOrNull(leaderId);

                    LOG.debug(
                            "New leader information for job {}. Address: {}, leader id: {}.",
                            consumerId,
                            leaderAddress,
                            jobMasterId);

                    if (leaderAddress == null || leaderAddress.isEmpty()) {
                        // the leader lost leadership but there is no other leader yet.
                        jobManagerLostLeadership = Optional.ofNullable(currentResourceConsumerId);
                        closeRpcConnection();
                    } else {
                        // check whether we are already connecting to this leader
                        if (Objects.equals(jobMasterId, currentResourceConsumerId)) {
                            LOG.debug(
                                    "Ongoing attempt to connect to leader of job {}. Ignoring duplicate leader information.",
                                    consumerId);
                        } else {
                            closeRpcConnection();
                            openRpcConnectionTo(leaderAddress, jobMasterId);
                        }
                    }
                }
            }

            // send callbacks outside of the lock scope
            jobManagerLostLeadership.ifPresent(
                    oldJobMasterId ->
                            consumerLeaderListener.resourceConsumerLostLeadership(consumerId, oldJobMasterId));
        }

        @GuardedBy("lock")
        private void openRpcConnectionTo(String leaderAddress, ConsumingMasterId jobMasterId) {
            Preconditions.checkState(
                    currentResourceConsumerId == null && rpcConnection == null,
                    "Cannot open a new rpc connection if the previous connection has not been closed.");

            currentResourceConsumerId = jobMasterId;
            rpcConnection =
                    new ResourceConsumerRegisteredRpcConnection(
                            LOG, leaderAddress, jobMasterId, rpcService.getScheduledExecutor());

            LOG.info(
                    "Try to register at job manager {} with leader id {}.",
                    leaderAddress,
                    jobMasterId.toUUID());
            rpcConnection.start();
        }

        @GuardedBy("lock")
        private void closeRpcConnection() {
            if (rpcConnection != null) {
                rpcConnection.close();
                rpcConnection = null;
                currentResourceConsumerId = null;
            }
        }

        @Override
        public void handleError(Exception exception) {
            if (stopped) {
                LOG.debug(
                        "{}'s leader retrieval listener reported an exception for job {}. "
                                + "However, the service is no longer running.",
                        DefaultConsumerLeaderService.class.getSimpleName(),
                        consumerId,
                        exception);
            } else {
                consumerLeaderListener.handleError(exception);
            }
        }

        /** Rpc connection for the resource consumer <--> resource provider connection. */
        private final class ResourceConsumerRegisteredRpcConnection
                extends RegisteredRpcConnection<
                		ConsumingMasterId,
                		ResourceConsumerGateway,
                        RCRPRegistrationSuccess,
                        RCRPRegistrationRejection> {

            ResourceConsumerRegisteredRpcConnection(
                    Logger log, String targetAddress, ConsumingMasterId jobMasterId, Executor executor) {
                super(log, targetAddress, jobMasterId, executor);
            }

            @Override
            protected RetryingRegistration<
            				ConsumingMasterId,
            				ResourceConsumerGateway,
                            RCRPRegistrationSuccess,
                            RCRPRegistrationRejection>
                    generateRegistration() {
                return new DefaultConsumerLeaderService.ResourceProviderRetryingRegistration(
                        LOG,
                        rpcService,
                        "ResourceConsumer",
                        ResourceConsumerGateway.class,
                        getTargetAddress(),
                        getTargetLeaderId(),
                        retryingRegistrationConfiguration,
                        consumerId,
                        ResourceProviderRegistrationInformation.create(
                                ownerAddress, ownLocation, taskManagerSession));
            }

            @Override
            protected void onRegistrationSuccess(RCRPRegistrationSuccess success) {
                runIfValidRegistrationAttemptOrElse(
                        () -> {
                            log.info(
                                    "Successful registration at job manager {} for job {}.",
                                    getTargetAddress(),
                                    consumerId);

                            consumerLeaderListener.resourceConsumerGainedLeadership(
                                    consumerId, getTargetGateway(), success);
                        },
                        () ->
                                log.debug(
                                        "Encountered obsolete JobManager registration success from {} with leader session ID {}.",
                                        getTargetAddress(),
                                        getTargetLeaderId()));
            }

            @Override
            protected void onRegistrationRejection(RCRPRegistrationRejection rejection) {
                runIfValidRegistrationAttemptOrElse(
                        () -> {
                            log.info(
                                    "Rejected registration at job manager {} for job {}.",
                                    getTargetAddress(),
                                    consumerId);

                            consumerLeaderListener.consumerManagerRejectedRegistration(
                                    consumerId, getTargetAddress(), rejection);
                        },
                        () ->
                                log.debug(
                                        "Encountered obsolete JobManager registration rejection {} from {} with leader session ID {}.",
                                        rejection,
                                        getTargetAddress(),
                                        getTargetLeaderId()));
            }

            @Override
            protected void onRegistrationFailure(Throwable failure) {
                runIfValidRegistrationAttemptOrElse(
                        () -> {
                            log.info(
                                    "Failed to register at job  manager {} for job {}.",
                                    getTargetAddress(),
                                    consumerId);
                            consumerLeaderListener.handleError(failure);
                        },
                        () ->
                                log.debug(
                                        "Obsolete JobManager registration failure from {} with leader session ID {}.",
                                        getTargetAddress(),
                                        getTargetLeaderId(),
                                        failure));
            }

            private void runIfValidRegistrationAttemptOrElse(
                    Runnable runIfValid, Runnable runIfInvalid) {
                if (Objects.equals(getTargetLeaderId(), getCurrentJobMasterId())) {
                    runIfValid.run();
                } else {
                    runIfInvalid.run();
                }
            }
        }
    }

    /** Retrying registration for the resource consumer <--> resource provider connection. */
    private static final class ResourceProviderRetryingRegistration
            extends RetryingRegistration<
                    ConsumingMasterId,
                    ResourceConsumerGateway,
                    RCRPRegistrationSuccess,
                    RCRPRegistrationRejection> {

        private final ConsumerID consumerId;

        private final ResourceProviderRegistrationInformation taskManagerRegistrationInformation;

        ResourceProviderRetryingRegistration(
                Logger log,
                RpcService rpcService,
                String targetName,
                Class<ResourceConsumerGateway> targetType,
                String targetAddress,
                ConsumingMasterId jobMasterId,
                RetryingRegistrationConfiguration retryingRegistrationConfiguration,
                ConsumerID consumerId,
                ResourceProviderRegistrationInformation taskManagerRegistrationInformation) {
            super(
                    log,
                    rpcService,
                    targetName,
                    targetType,
                    targetAddress,
                    jobMasterId,
                    retryingRegistrationConfiguration);

            this.consumerId = Preconditions.checkNotNull(consumerId);
            this.taskManagerRegistrationInformation = taskManagerRegistrationInformation;
        }

        @Override
        protected CompletableFuture<RegistrationResponse> invokeRegistration(
                ResourceConsumerGateway gateway, ConsumingMasterId fencingToken, long timeoutMillis) {
            return gateway.registerResourceProvider(
                    consumerId, taskManagerRegistrationInformation, Time.milliseconds(timeoutMillis));
        }
    }

    /** Internal state of the service. */
    private enum State {
        CREATED,
        STARTED,
        STOPPED
    }

    // -----------------------------------------------------------
    // Testing methods
    // -----------------------------------------------------------

    /**
     * Check whether the service monitors the given job.
     *
     * @param jobId identifying the job
     * @return True if the given job is monitored; otherwise false
     */
    @Override
    @VisibleForTesting
    public boolean containsJob(ConsumerID jobId) {
        Preconditions.checkState(
                DefaultConsumerLeaderService.State.STARTED == state,
                "The service is currently not running.");

        return consumerLeaderServices.containsKey(jobId);
    }
}
